{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
    "os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'prepare/mesolitica-tpu.json'\n",
    "b2_application_key_id = os.environ['b2_application_key_id']\n",
    "b2_application_key = os.environ['b2_application_key']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import storage\n",
    "client = storage.Client()\n",
    "bucket = client.bucket('mesolitica-tpu-general')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "best = '1050000'\n",
    "directory = 't5-super-super-tiny-segmentation'\n",
    "!rm -rf output out {directory}\n",
    "!mkdir {directory}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = best\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.data-00000-of-00002')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.data-00000-of-00002')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.data-00001-of-00002')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.data-00001-of-00002')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.index')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.index')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.meta')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.meta')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/checkpoint')\n",
    "blob.download_to_filename(f'{directory}/checkpoint')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/operative_config.gin')\n",
    "blob.download_to_filename(f'{directory}/operative_config.gin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'{directory}/checkpoint', 'w') as fopen:\n",
    "    fopen.write(f'model_checkpoint_path: \"model.ckpt-{model}\"')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from b2sdk.v1 import *\n",
    "info = InMemoryAccountInfo()\n",
    "b2_api = B2Api(info)\n",
    "application_key_id = b2_application_key_id\n",
    "application_key = b2_application_key\n",
    "b2_api.authorize_account(\"production\", application_key_id, application_key)\n",
    "file_info = {'how': 'good-file'}\n",
    "b2_bucket = b2_api.get_bucket_by_name('malaya-model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7fa61dfd4c18>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tar = 't5-super-super-tiny-segmentation-2021-09-10.tar.gz'\n",
    "os.system(f'tar -czvf {tar} {directory}')\n",
    "outPutname = f'finetuned/{tar}'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=tar,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'rm {tar}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import t5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = t5.models.MtfModel(\n",
    "    model_dir=directory,\n",
    "    tpu=None,\n",
    "    tpu_topology=None,\n",
    "    model_parallelism=1,\n",
    "    batch_size=1,\n",
    "    sequence_length={\"inputs\": 256, \"targets\": 256},\n",
    "    learning_rate_schedule=0.003,\n",
    "    save_checkpoints_steps=5000,\n",
    "    keep_checkpoint_max=3,\n",
    "    iterations_per_loop=100,\n",
    "    mesh_shape=\"model:1,batch:1\", \n",
    "    mesh_devices=[\"cpu:0\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf output/*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': 't5-super-super-tiny-segmentation', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': None, '_session_config': allow_soft_placement: true\n",
      "isolate_session_state: true\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_service': None, '_cluster_spec': None, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=100, num_shards=None, num_cores_per_replica=1, per_host_input_for_training=4, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None, eval_training_input_configuration=2, experimental_host_call_every_n_steps=1), '_cluster': <tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver.TPUClusterResolver object at 0x7fa5127a2588>}\n",
      "INFO:tensorflow:_TPUContext: eval_on_tpu True\n",
      "WARNING:tensorflow:eval_on_tpu ignored because use_tpu is False.\n"
     ]
    }
   ],
   "source": [
    "import gin\n",
    "\n",
    "from t5.data import sentencepiece_vocabulary\n",
    "\n",
    "DEFAULT_SPM_PATH = 'prepare/sp10m.cased.ms-en.model'\n",
    "DEFAULT_EXTRA_IDS = 100\n",
    "model_dir = directory\n",
    "\n",
    "def get_default_vocabulary():\n",
    "    return sentencepiece_vocabulary.SentencePieceVocabulary(\n",
    "      DEFAULT_SPM_PATH, DEFAULT_EXTRA_IDS)\n",
    "\n",
    "with gin.unlock_config():\n",
    "    gin.parse_config_file(t5.models.mtf_model._operative_config_path(model_dir))\n",
    "    gin.bind_parameter(\"Bitransformer.decode.beam_size\", 1)\n",
    "    gin.bind_parameter(\"Bitransformer.decode.temperature\", 0)\n",
    "    gin.bind_parameter(\"utils.get_variable_dtype.slice_dtype\", \"float32\")\n",
    "    gin.bind_parameter(\n",
    "        \"utils.get_variable_dtype.activation_dtype\", \"float32\")\n",
    "    \n",
    "vocabulary = t5.data.SentencePieceVocabulary(DEFAULT_SPM_PATH)\n",
    "estimator = model.estimator(vocabulary, disable_tpu=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1050000,\n",
       " 'model.ckpt-1050000',\n",
       " 't5-super-super-tiny-segmentation/model.ckpt-1050000')"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "checkpoint_step = t5.models.mtf_model._get_latest_checkpoint_from_dir(model_dir)\n",
    "model_ckpt = \"model.ckpt-\" + str(checkpoint_step)\n",
    "checkpoint_path = os.path.join(model_dir, model_ckpt)\n",
    "checkpoint_step, model_ckpt, checkpoint_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1630: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Running infer on CPU\n",
      "INFO:tensorflow:feature inputs : Tensor(\"Reshape:0\", shape=(1, 256), dtype=int32)\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/mesh_tensorflow/transformer/utils.py:427: Print (from tensorflow.python.ops.logging_ops) is deprecated and will be removed after 2018-08-20.\n",
      "Instructions for updating:\n",
      "Use tf.print instead of tf.Print. Note that tf.print returns a no-output operator that directly prints the output. Outside of defuns or eager mode, this operator will not be executed unless it is directly specified in session.run or used as a control dependency for other operators. This is only a concern in graph mode. Below is an example of how to ensure tf.print executes in graph mode:\n",
      "\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable encoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/k                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/o                  size 49152        slice_size 49152        Shape[heads=384, d_model=128]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/q                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias size 192          slice_size 192          Shape[heads=6, buckets=32]                                  \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/v                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/k                size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/o                size 49152        slice_size 49152        Shape[heads=384, d_model=128]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/q                size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/v                size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/DenseReluDense/wi/kernel         size 65536        slice_size 65536        Shape[d_model=128, d_ff=512]                                \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/DenseReluDense/wo/kernel         size 65536        slice_size 65536        Shape[d_ff=512, d_model=128]                                \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/k                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/o                  size 49152        slice_size 49152        Shape[heads=384, d_model=128]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/q                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/v                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/k                size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/o                size 49152        slice_size 49152        Shape[heads=384, d_model=128]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/q                size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/v                size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/DenseReluDense/wi/kernel         size 65536        slice_size 65536        Shape[d_model=128, d_ff=512]                                \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/DenseReluDense/wo/kernel         size 65536        slice_size 65536        Shape[d_ff=512, d_model=128]                                \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable decoder/final_layer_norm/scale                               size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/k                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/o                  size 49152        slice_size 49152        Shape[heads=384, d_model=128]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/q                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/relative_attention_bias size 192          slice_size 192          Shape[heads=6, buckets=32]                                  \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/v                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/DenseReluDense/wi/kernel         size 65536        slice_size 65536        Shape[d_model=128, d_ff=512]                                \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/DenseReluDense/wo/kernel         size 65536        slice_size 65536        Shape[d_ff=512, d_model=128]                                \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/k                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/o                  size 49152        slice_size 49152        Shape[heads=384, d_model=128]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/q                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/v                  size 49152        slice_size 49152        Shape[d_model=128, heads=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/DenseReluDense/wi/kernel         size 65536        slice_size 65536        Shape[d_model=128, d_ff=512]                                \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/DenseReluDense/wo/kernel         size 65536        slice_size 65536        Shape[d_ff=512, d_model=128]                                \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/layer_norm/scale                 size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable encoder/final_layer_norm/scale                               size 128          slice_size 128          Shape[d_model=128]                                          \n",
      "INFO:tensorflow:Variable shared/embedding                                             size 4112384      slice_size 4112384      Shape[vocab=32128, d_model=128]                             \n",
      "INFO:tensorflow:Trainable Variables            count: 47      Total size: 5818240          Total slice_size: 5818240        \n",
      "INFO:tensorflow:All Variables                  count: 47      Total size: 5818240          Total slice_size: 5818240        \n",
      "INFO:tensorflow:Counters:\n",
      "einsum: 3.96e+09\n",
      "einsum_unique: 3.96e+09\n",
      "output: 7.53e+07\n",
      " output/AddOperation: 9.64e+06\n",
      " output/BinaryOpWithBroadcasting: 6.57e+05\n",
      " output/Constant: 3.93e+05\n",
      " output/EinsumOperation: 1.76e+07\n",
      " output/ImportOperation: 277\n",
      " output/MinMaxOperation: 3.93e+05\n",
      " output/OneHotOperation: 2.48e+07\n",
      " output/RangeOperation: 512\n",
      " output/ReduceOperation: 2.15e+04\n",
      " output/ReshapeOperation: 3.12e+06\n",
      " output/ScalarAddOperation: 5.27e+05\n",
      " output/ScalarMultiplyOperation: 1.22e+06\n",
      " output/ShiftOperation: 256\n",
      " output/SlicewiseOperation: 8.29e+06\n",
      " output/StopGradient: 2.36e+06\n",
      " output/Variable: 5.82e+06\n",
      " output/WhileLoopOperation: 3.93e+05\n",
      "output_unique: 7.53e+07\n",
      " output_unique/AddOperation: 9.64e+06\n",
      " output_unique/BinaryOpWithBroadcasting: 6.57e+05\n",
      " output_unique/Constant: 3.93e+05\n",
      " output_unique/EinsumOperation: 1.76e+07\n",
      " output_unique/ImportOperation: 277\n",
      " output_unique/MinMaxOperation: 3.93e+05\n",
      " output_unique/OneHotOperation: 2.48e+07\n",
      " output_unique/RangeOperation: 512\n",
      " output_unique/ReduceOperation: 2.15e+04\n",
      " output_unique/ReshapeOperation: 3.12e+06\n",
      " output_unique/ScalarAddOperation: 5.27e+05\n",
      " output_unique/ScalarMultiplyOperation: 1.22e+06\n",
      " output_unique/ShiftOperation: 256\n",
      " output_unique/SlicewiseOperation: 8.29e+06\n",
      " output_unique/StopGradient: 2.36e+06\n",
      " output_unique/Variable: 5.82e+06\n",
      " output_unique/WhileLoopOperation: 3.93e+05\n",
      "variables: 5.82e+06\n",
      " variables/trainable: 5.82e+06\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/ops/array_ops.py:1475: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/saved_model/signature_def_utils_impl.py:201: build_tensor_info (from tensorflow.python.saved_model.utils_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.utils.build_tensor_info or tf.compat.v1.saved_model.build_tensor_info.\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Classify: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Regress: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Predict: ['serving_default']\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Train: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Eval: None\n",
      "INFO:tensorflow:Restoring parameters from t5-super-super-tiny-segmentation/model.ckpt-1050000\n",
      "INFO:tensorflow:Assets added to graph.\n",
      "INFO:tensorflow:No assets to write.\n",
      "INFO:tensorflow:SavedModel written to: output/temp-b'1631723970'/saved_model.pb\n"
     ]
    }
   ],
   "source": [
    "from mesh_tensorflow.transformer import dataset as transformer_dataset\n",
    "\n",
    "def serving_input_fn():\n",
    "    inputs = tf.placeholder(\n",
    "            dtype=tf.string,\n",
    "            shape=[None],\n",
    "            name=\"inputs\")\n",
    "\n",
    "    batch_size = tf.shape(inputs)[0]\n",
    "    padded_inputs = tf.pad(inputs, [(0, tf.mod(-tf.size(inputs), batch_size))])\n",
    "    dataset = tf.data.Dataset.from_tensor_slices(padded_inputs)\n",
    "    dataset = dataset.map(lambda x: {\"inputs\": x})\n",
    "    dataset = transformer_dataset.encode_all_features(dataset, vocabulary)\n",
    "    dataset = transformer_dataset.pack_or_pad(\n",
    "        dataset=dataset,\n",
    "        length=model._sequence_length,\n",
    "        pack=False,\n",
    "        feature_keys=[\"inputs\"]\n",
    "    )\n",
    "    dataset = dataset.batch(tf.cast(batch_size, tf.int64))\n",
    "    features = tf.data.experimental.get_single_element(dataset)\n",
    "    return tf.estimator.export.ServingInputReceiver(\n",
    "        features=features, receiver_tensors=inputs)\n",
    "\n",
    "out = estimator.export_saved_model('output', serving_input_fn, checkpoint_path=checkpoint_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-17-5b89b6a20c22>:7: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.\n",
      "INFO:tensorflow:Restoring parameters from output/1631723970/variables/variables\n"
     ]
    }
   ],
   "source": [
    "config = tf.ConfigProto()\n",
    "config.allow_soft_placement = True\n",
    "sess = tf.Session(config = config)\n",
    "meta_graph_def = tf.saved_model.loader.load(\n",
    "        sess,\n",
    "        [tf.saved_model.tag_constants.SERVING],\n",
    "        out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'super-super-tiny-segmentation/model.ckpt'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.save(sess, 'super-super-tiny-segmentation/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "strings = [\n",
    "    n.name\n",
    "    for n in tf.get_default_graph().as_graph_def().node\n",
    "    if ('encoder' in n.op\n",
    "    or 'decoder' in n.name\n",
    "    or 'shared' in n.name\n",
    "    or 'inputs' in n.name\n",
    "    or 'output' in n.name\n",
    "    or 'SentenceTokenizer' in n.name\n",
    "    or 'self/Softmax' in n.name)\n",
    "    and 'adam' not in n.name\n",
    "    and 'Assign' not in n.name\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze_graph(model_dir, output_node_names):\n",
    "\n",
    "    if not tf.gfile.Exists(model_dir):\n",
    "        raise AssertionError(\n",
    "            \"Export directory doesn't exists. Please specify an export \"\n",
    "            'directory: %s' % model_dir\n",
    "        )\n",
    "\n",
    "    checkpoint = tf.train.get_checkpoint_state(model_dir)\n",
    "    input_checkpoint = checkpoint.model_checkpoint_path\n",
    "\n",
    "    absolute_model_dir = '/'.join(input_checkpoint.split('/')[:-1])\n",
    "    output_graph = absolute_model_dir + '/frozen_model.pb'\n",
    "    clear_devices = True\n",
    "    with tf.Session(graph = tf.Graph()) as sess:\n",
    "        saver = tf.train.import_meta_graph(\n",
    "            input_checkpoint + '.meta', clear_devices = clear_devices\n",
    "        )\n",
    "        saver.restore(sess, input_checkpoint)\n",
    "        output_graph_def = tf.graph_util.convert_variables_to_constants(\n",
    "            sess,\n",
    "            tf.get_default_graph().as_graph_def(),\n",
    "            output_node_names,\n",
    "        )\n",
    "        with tf.gfile.GFile(output_graph, 'wb') as f:\n",
    "            f.write(output_graph_def.SerializeToString())\n",
    "        print('%d ops in the final graph.' % len(output_graph_def.node))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from super-super-tiny-segmentation/model.ckpt\n",
      "WARNING:tensorflow:From <ipython-input-20-504c79665720>:23: convert_variables_to_constants (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.convert_variables_to_constants`\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/framework/graph_util_impl.py:277: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.extract_sub_graph`\n",
      "INFO:tensorflow:Froze 76 variables.\n",
      "INFO:tensorflow:Converted 76 variables to const ops.\n",
      "2616 ops in the final graph.\n"
     ]
    }
   ],
   "source": [
    "freeze_graph('super-super-tiny-segmentation', strings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import struct\n",
    "\n",
    "unknown = b'\\xff\\xff\\xff\\xff'\n",
    "\n",
    "def load_graph(frozen_graph_filename):\n",
    "    with tf.gfile.GFile(frozen_graph_filename, 'rb') as f:\n",
    "        graph_def = tf.GraphDef()\n",
    "        graph_def.ParseFromString(f.read())\n",
    "        \n",
    "    for node in graph_def.node:\n",
    "        \n",
    "        if node.op == 'RefSwitch':\n",
    "          node.op = 'Switch'\n",
    "          for index in xrange(len(node.input)):\n",
    "            if 'moving_' in node.input[index]:\n",
    "              node.input[index] = node.input[index] + '/read'\n",
    "        elif node.op == 'AssignSub':\n",
    "          node.op = 'Sub'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "        elif node.op == 'AssignAdd':\n",
    "          node.op = 'Add'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "        elif node.op == 'Assign':\n",
    "          node.op = 'Identity'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "          if 'validate_shape' in node.attr: del node.attr['validate_shape']\n",
    "          if len(node.input) == 2:\n",
    "            node.input[0] = node.input[1]\n",
    "            del node.input[1]\n",
    "            \n",
    "        if 'Reshape/shape' in node.name or 'Reshape_1/shape' in node.name:\n",
    "            b = node.attr['value'].tensor.tensor_content\n",
    "            arr_int = [int.from_bytes(b[i:i + 4], 'little') for i in range(0, len(b), 4)]\n",
    "            if len(arr_int):\n",
    "                arr_byte = [unknown] + [struct.pack('<i', i) for i in arr_int[1:]]\n",
    "                arr_byte = b''.join(arr_byte)\n",
    "                node.attr['value'].tensor.tensor_content = arr_byte\n",
    "            \n",
    "            if len(node.attr['value'].tensor.int_val):\n",
    "                node.attr['value'].tensor.int_val[0] = -1\n",
    "    \n",
    "    with tf.Graph().as_default() as graph:\n",
    "        tf.import_graph_def(graph_def)\n",
    "    return graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = load_graph('super-super-tiny-segmentation/frozen_model.pb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor 'import/inputs:0' shape=(?,) dtype=string>,\n",
       " <tf.Tensor 'import/SelectV2_3:0' shape=(?, 256) dtype=int32>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i = g.get_tensor_by_name('import/inputs:0')\n",
    "o = g.get_tensor_by_name('import/SelectV2_3:0')\n",
    "i, o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sess = tf.Session(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sentencepiece as spm\n",
    "sp_model = spm.SentencePieceProcessor()\n",
    "sp_model.Load(DEFAULT_SPM_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "string1 = 'krajaanpatutbagipencen awal skt kpd warga emas supaya emosi'\n",
    "string2 = 'Husein ska mkn aym dkat kampng Jawa'\n",
    "string3 = 'Melayumalasninarration dia samaje macam men are trash.sakittengok'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['segmentasi: krajaanpatutbagipencen awal skt kpd warga emas supaya emosi',\n",
       " 'segmentasi: Husein ska mkn aym dkat kampng Jawa',\n",
       " 'segmentasi: Melayumalasninarration dia samaje macam men are trash.sakittengok']"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "strings = [string1, string2, string3]\n",
    "[f'segmentasi: {s}' for s in strings]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.08 s, sys: 731 ms, total: 3.81 s\n",
      "Wall time: 1.73 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3, 256)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "o_ = test_sess.run(o, feed_dict = {i: [f'segmentasi: {s}' for s in strings]})\n",
    "o_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 krajaan patut bagi pencen awal skt kpd warga emas supaya emosi\n",
      "1 Husein ska mkn aym dkat kampng Jawa\n",
      "2 Melayu malas ninarration dia sama je macam men are trash. sakit tengok\n"
     ]
    }
   ],
   "source": [
    "for k in range(len(o_)):\n",
    "    print(k, sp_model.DecodeIds(o_[k].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.tools.graph_transforms import TransformGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms = ['add_default_attributes',\n",
    "             'remove_nodes(op=Identity, op=CheckNumerics)',\n",
    "             'fold_batch_norms',\n",
    "             'fold_old_batch_norms',\n",
    "             'quantize_weights(minimum_size=1536000)',\n",
    "             #'quantize_weights(fallback_min=-10240, fallback_max=10240)',\n",
    "             'strip_unused_nodes',\n",
    "             'sort_by_execution_order']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-33-bd88caf3ba4d>:3: FastGFile.__init__ (from tensorflow.python.platform.gfile) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.gfile.GFile.\n"
     ]
    }
   ],
   "source": [
    "pb = 'super-super-tiny-segmentation/frozen_model.pb'\n",
    "input_graph_def = tf.GraphDef()\n",
    "with tf.gfile.FastGFile(pb, 'rb') as f:\n",
    "    input_graph_def.ParseFromString(f.read())\n",
    "    \n",
    "transformed_graph_def = TransformGraph(input_graph_def, \n",
    "       ['inputs'],\n",
    "       ['SelectV2_3'], transforms)\n",
    "\n",
    "with tf.gfile.GFile(f'{pb}.quantized', 'wb') as f:\n",
    "    f.write(transformed_graph_def.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor 'import/inputs:0' shape=(?,) dtype=string>,\n",
       " <tf.Tensor 'import/SelectV2_3:0' shape=(?, 256) dtype=int32>)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g = load_graph('super-super-tiny-segmentation/frozen_model.pb.quantized')\n",
    "i = g.get_tensor_by_name('import/inputs:0')\n",
    "o = g.get_tensor_by_name('import/SelectV2_3:0')\n",
    "i, o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sess = tf.InteractiveSession(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7fa4bc24c0b8>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file = 'super-super-tiny-segmentation/frozen_model.pb.quantized'\n",
    "outPutname = 'segmentation/super-super-tiny-t5-quantized/model.pb'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=file,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7fa61ff3bd30>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file = 'super-super-tiny-segmentation/frozen_model.pb'\n",
    "outPutname = 'segmentation/super-super-tiny-t5/model.pb'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=file,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
